#!/bin/bash

if [ "$EUID" -ne 0 ]
then
    echo "Must run as root or with sudo."
    exit 1
fi

TARBALL=${1}
CONFIG=${2}
DRIVESIZE=${3}
DEVICE=`losetup -f`
DEVICE_BOOT=${DEVICE}p1
DEVICE_ROOT=${DEVICE}p2

BOOT_MNT=/mnt/boot
ROOT_MNT=/mnt/root

if [[ -z "$TARBALL" ]]
then
    echo "Usage: ./mkeosimg system_image.tar <optional 2nd parameter: /path/to/configfile.tar.gz> \
    <optional 3rd parameter: desired image size in GB>"
    echo "If you want to provide an image size but no config, use \"-\" without the quotes as the second paramenter."
    exit 1
fi

mkdir -p $BOOT_MNT
mkdir -p $ROOT_MNT

# unmount only if already mounted
grep -qs "$BOOT_MNT" /proc/mounts && umount $BOOT_MNT
grep -qs "$ROOT_MNT" /proc/mounts && umount $ROOT_MNT

DIR=`mktemp -d`

function cleanup {
    if [ -d "$DIR" ]
    then
        rm -rf "$DIR"
    fi
    if [ -d "$BOOT_MNT" ]
    then
        rm -rf "$BOOT_MNT"
    fi
    if [ -d "$ROOT_MNT" ]
    then
        rm -rf "$ROOT_MNT"
    fi
}
trap cleanup EXIT

####### Check if config was supplied as a parameter and if it exists as a file
if [ ! -z $CONFIG ] && [ -f $CONFIG ]
then
    CONFIGURED_SUFFIX="-configured"
else
    CONFIGURED_SUFFIX=""
fi

IMG=${TARBALL%.tar}${CONFIGURED_SUFFIX}.img
rm -f $IMG.gz

# If no drive size was supplied as a parameter, default to 2 GB
if [ -z $DRIVESIZE ]
then
    DRIVESIZE=2
fi

echo "Creating $IMG with a size of $DRIVESIZE GB"
let "SIZE = $DRIVESIZE * 1024"
dd if=/dev/zero of=$IMG bs=1M count=0 seek=$SIZE
losetup $DEVICE $IMG
parted --script $DEVICE mktable msdos
parted --script $DEVICE mkpart primary fat32 1 150MB
parted --script $DEVICE mkpart primary ext3 150MB 100%

echo "Formatting $IMG"
mkfs.msdos -F 32 $DEVICE_BOOT
mkfs.ext3 -q $DEVICE_ROOT

mount $DEVICE_BOOT $BOOT_MNT
mount $DEVICE_ROOT $ROOT_MNT

echo "Unpacking EdgeOS release image"
tar xf $TARBALL -C $DIR

####### Kernel
echo "Verifying EdgeOS kernel"
if [ `md5sum $DIR/vmlinux.tmp | awk -F ' ' '{print $1}'` != \
     `cat $DIR/vmlinux.tmp.md5` ]; then
    echo "Kernel in image is corrupted! Check your image and start over."
    exit 1
fi

echo "Copying EdgeOS kernel to boot partition"
cp $DIR/vmlinux.tmp $BOOT_MNT/vmlinux.64
cp $DIR/vmlinux.tmp.md5 $BOOT_MNT/vmlinux.64.md5

####### System
echo "Verifying EdgeOS system image"
if [ `md5sum $DIR/squashfs.tmp | awk -F ' ' '{print $1}'` != \
     `cat $DIR/squashfs.tmp.md5` ]; then
    echo "System in image is corrupted! Check your image and start over."
    exit 1
fi

echo "Copying EdgeOS system image to root partition"
mv $DIR/squashfs.tmp $ROOT_MNT/squashfs.img
mv $DIR/squashfs.tmp.md5 $ROOT_MNT/squashfs.img.md5

echo "Copying version file to the root partition"
mv $DIR/version.tmp $ROOT_MNT/version

# Writable data dir
echo "Creating EdgeOS writable data directories"
mkdir $ROOT_MNT/w
mkdir $ROOT_MNT/www

chown 33:0 $ROOT_MNT/www

####### Extract config
if [ ! -z $CONFIG ] && [ -f $CONFIG ]
then
    tar xf $CONFIG -C $ROOT_MNT/w
fi

# These need to happen before the losetup -d, or else the changes won't be
# written to disk
umount $ROOT_MNT
umount $BOOT_MNT

# Unhook img
losetup -d $DEVICE

echo "Done."
